# coding=utf-8

# author: zhangshu
# 本工具用来实现旋转矩阵、欧拉角、四元数之间的相互转换
# 用于TF的外参配置等功能

import numpy as np
from scipy.spatial.transform import Rotation
import rospy
 
import math
import numpy as np
import tf

rospy.init_node("py_tools")
 
def eulerAngles2rotationMat(theta, format='degree'):
    if format is 'degree':
        theta = [i * math.pi / 180.0 for i in theta]
 
    R_x = np.array([[1, 0, 0],
                    [0, math.cos(theta[0]), -math.sin(theta[0])],
                    [0, math.sin(theta[0]), math.cos(theta[0])]
                    ])
 
    R_y = np.array([[math.cos(theta[1]), 0, math.sin(theta[1])],
                    [0, 1, 0],
                    [-math.sin(theta[1]), 0, math.cos(theta[1])]
                    ])
 
    R_z = np.array([[math.cos(theta[2]), -math.sin(theta[2]), 0],
                    [math.sin(theta[2]), math.cos(theta[2]), 0],
                    [0, 0, 1]
                    ])
    R = np.dot(R_z, np.dot(R_y, R_x))
    return R

# 欧拉角转旋转矩阵 3*3
print(eulerAngles2rotationMat([0,0,-90]))

roll = 0
pitch = 0.0
yaw = -math.pi/2.0


# 4*4 欧拉角转旋转矩阵
matrix = tf.transformations.euler_matrix(roll,pitch,yaw)
print(matrix)


# 旋转矩阵转欧拉角
matrix = [[1, 0, 0],
          [0, 1, 0],
          [0, 0, 1]]
(roll,pitch,yaw) = tf.transformations.euler_from_matrix(matrix)
print("roll={}, pitch={}, yaw={}".format(roll, pitch, yaw))

# euler to quantation
from tf.transformations import quaternion_from_euler
from geometry_msgs.msg import Pose

pos = Pose()
DE2RA = 3.1415926 / 180

# 欧拉角转四元数
roll = 0
pitch = 0
yaw = -90
q = quaternion_from_euler(roll * DE2RA, pitch * DE2RA, yaw * DE2RA)
print("qx={},qy={},qz={},qw={}".format(q[0],q[1],q[2],q[3]))


# q=[0, 0, 0.7071, 0.7071]
# 四元数转欧拉角
(r, p, y) = tf.transformations.euler_from_quaternion(q)
print("roll={}, pitch={}, yaw={}".format(r, p, y))

